from utils.postgis import *
from shapely.geometry import LineString, Point, Polygon
import matplotlib.pyplot as plt


def is_same_point(coord1, coord2, tolerance=CLUSTER_TOLERANCE):
    if np.abs(coord1[0] - coord2[0]) <= tolerance and np.abs(coord1[1] - coord2[1]) <= tolerance:
        return True
    else:
        return False


def is_equal_point(coord1, coord2):
    if coord1[0] == coord2[0] and coord1[1] == coord2[1]:
        return True
    else:
        return False


def join(geo1, geo2):
    g2 = list(geo2.coords)
    g1 = list(geo1.coords)
    if is_same_point(geo1.coords[-1], geo2.coords[0]):
        g1.extend(g2[1:])
        g = LineString(g1)
        fun_val = FunEnum.JOIN_1
    elif is_same_point(geo2.coords[-1], geo1.coords[0]):
        g2.extend(g1[1:])
        g = LineString(g2)
        fun_val = FunEnum.JOIN_2
    elif is_same_point(geo1.coords[0], geo2.coords[0]):
        g1.reverse()
        g1.extend(g2[1:])
        g = LineString(g1)
        fun_val = FunEnum.JOIN_3
    elif is_same_point(geo1.coords[-1], geo2.coords[-1]):
        g2 = list(geo2.coords)
        g1 = list(geo1.coords)
        g2.reverse()
        g1.extend(g2[1:])
        g = LineString(g1)
        fun_val = FunEnum.JOIN_4
    else:  # 不是伪结点，是相交
        raise
    return g, fun_val


def extend(geom1, geom2, extend_threshold):
    # 找最近点
    geom1_list = list(geom1.coords)
    geom2_list = list(geom2.coords)

    st_1 = geom1_list[0]
    ed_1 = geom1_list[-1]
    st_2 = geom2_list[0]
    ed_2 = geom2_list[-1]

    dist_list = [cal_dist(st_1, st_2), cal_dist(st_1, ed_2), cal_dist(ed_1, st_2), cal_dist(ed_1, ed_2)]
    min_dist = np.min(dist_list)
    min_ind = dist_list.index(min_dist)

    if min_dist > extend_threshold:
        raise ExtendException(min_dist, extend_threshold)  # 可能是延伸到其他道路的中段了，也可能是其他错误

    if min_ind == 0:
        geom1_list.reverse()
        new_geom_list = geom1_list + geom2_list
    elif min_ind == 1:
        new_geom_list = geom2_list + geom1_list
    elif min_ind == 2:
        new_geom_list = geom1_list + geom2_list
    else:
        geom2_list.reverse()
        new_geom_list = geom1_list + geom2_list

    new_geom = LineString(new_geom_list)
    return new_geom


class ExtendException(Exception):
    def __init__(self, distance, extend_threshold):
        super().__init__()
        self.distance = distance
        self.extend_threshold = extend_threshold

    def __str__(self):
        return f"extend exceeds threshold: {self.distance} > {self.extend_threshold}"


class SnapException(Exception):
    def __init__(self, pt, ls, snap_threshold):
        super().__init__()
        self.pt = pt
        self.ls = ls
        self.snap_threshold = snap_threshold

    def __str__(self):
        return f"snap exceeds under threshold({self.snap_threshold}): {self.pt}, {self.ls.wkt}"


def project_pt_to_polyline(pt_coord, line_coord_list):
    """

    Args:
        pt_coord:
        line_coord_list:

    Returns: min_proj_dist, line_proj_status

    """
    pt_num = len(line_coord_list)
    min_proj_dist = np.inf
    proj_pt = (np.inf, np.inf)
    pre_node = -1  # 投影点的前一个结点
    for i in range(pt_num - 1):
        proj_result = project(pt_coord, line_coord_list[i], line_coord_list[i+1])
        proj_dist = proj_result["proj_dist"]
        proj_pt_temp = proj_result["proj_pt"]
        if proj_dist < min_proj_dist:
            min_proj_dist = proj_dist
            pre_node = i
            proj_pt = proj_pt_temp
    # proj_status: 0:投影到线段上 1:投影到起点 2:投影到终点
    if is_same_point(proj_pt, line_coord_list[0]):
        proj_status = 1
    elif is_same_point(proj_pt, line_coord_list[-1]):
        proj_status = 2
    else:
        proj_status = 0
    return {
        "proj_dist": min_proj_dist,
        "proj_status": proj_status,
        "pre_node": pre_node,
        "proj_pt": proj_pt
    }


def project(pt_coord: tuple, line_st: tuple, line_ed: tuple):
    """
    https://www.cnblogs.com/flyinggod/p/9359534.html

    Args:
        pt_coord:
        line_st:
        line_ed:

    Returns: px, py, cal_dist(x, y, px, py), projection_status

    """
    x, y = pt_coord
    x1, y1 = line_st
    x2, y2 = line_ed
    cross = (x2 - x1) * (x - x1) + (y2 - y1) * (y - y1)
    if cross <= 0:  # 投影到x1,y1上
        projection_status = 1
        px, py = x1, y1
    else:
        d2 = (x1 - x2) ** 2 + (y1 - y2) ** 2
        if cross >= d2:  # 投影到x2,y2上
            projection_status = 2
            px, py = x2, y2
        else:  # 投影到中间
            projection_status = 0
            r = cross / d2
            px = x1 + (x2 - x1) * r
            py = y1 + (y2 - y1) * r
    return {
        "proj_dist": cal_dist(pt_coord, (px, py)),
        "proj_pt": (px, py),
        "proj_status": projection_status
    }


def cal_dist(coord1, coord2):
    x1, y1 = coord1
    x2, y2 = coord2
    return np.sqrt((x1 - x2) ** 2 + (y1 - y2) ** 2)


def cal_dist_end(pt, line_pt_list):
    """
    计算pt到line两个端点的距离

    Args:
        pt:
        line_pt_list:

    Returns:
    [pt到st距离，pt到ed距离]

    """
    dist_st = cal_dist(pt, line_pt_list[0])
    dist_ed = cal_dist(pt, line_pt_list[-1])
    return [dist_st, dist_ed]


def plot_linestring(geom, c='b', alpha=1):
    x, y = geom.xy
    plt.plot(x, y, c=c, alpha=alpha)
    plt.axis('equal')


def split_line_at_pt(line_pt_list, pre_node, split_pt: tuple):
    """
    输入line的排列：..., pre_node, split_pt, pre_node+1, ...

    打断为：
        geom1：..., pre_node, split_pt
        geom2：split_pt, pre_node+1, ...

    特殊情况：
    1. split_pt与端点重合，不打断，返回一条线段
    2. split_pt与某结点重合，则将结点直接替换为split_pt，不重复插入。

    Args:
        line_pt_list:
        pre_node:
        split_pt:

    Returns: 列表，长度为1或2.

    """
    # 分割点与起点重合
    if is_same_point(split_pt, line_pt_list[0]):
        geom1_pt_list = []
        geom2_pt_list = line_pt_list
    # 分割点与终点重合
    elif is_same_point(split_pt, line_pt_list[-1]):
        geom1_pt_list = line_pt_list
        geom2_pt_list = []
    # 分割点与pre_node重合
    elif is_same_point(split_pt, line_pt_list[pre_node]):
        geom1_pt_list = line_pt_list[:pre_node] + [split_pt]  # 捕捉到split_coord上
        geom2_pt_list = [split_pt] + line_pt_list[pre_node + 1:]
    # 分割点与pre_node+1重合
    elif is_same_point(split_pt, line_pt_list[pre_node + 1]):
        geom1_pt_list = line_pt_list[:pre_node + 1] + [split_pt]
        geom2_pt_list = [split_pt] + line_pt_list[pre_node + 2:]
    # 分割点不在结点上
    else:
        geom1_pt_list = line_pt_list[:pre_node + 1] + [split_pt]
        geom2_pt_list = [split_pt] + line_pt_list[pre_node + 1:]

    split_line_list = []
    if len(geom1_pt_list) != 0:
        split_line_list.append(geom1_pt_list)
    if len(geom2_pt_list) != 0:
        split_line_list.append(geom2_pt_list)
    return split_line_list


def split_line_by_project(line_pt_list: list, pt: tuple):
    """
    用任意一点pt，在投影点上打断线段line_pt_list
    Args:
        line_pt_list:
        pt:

    Returns:

    """
    proj_result = project_pt_to_polyline(pt, line_pt_list)
    pre_node = proj_result["pre_node"]
    proj_node = proj_result["proj_pt"]
    split_list = split_line_at_pt(line_pt_list, pre_node, proj_node)
    return split_list


class PreprocessingTool:
    def __init__(self, table):
        self.table = table
        self.log = []
        self.pk_list = []

    def set_pk_list(self, pk_list):
        self.pk_list = pk_list


def get_substring(linestring, st_coord, ed_coord):
    start_dist = linestring.project(Point(st_coord))
    end_dist = linestring.project(Point(ed_coord))
    start_dist, end_dist = sorted([start_dist, end_dist])
    sub = shapely_substring(linestring, start_dist=start_dist, end_dist=end_dist)
    return sub


def shapely_substring(geom, start_dist, end_dist, normalized=False):
    """Return a line segment between specified distances along a LineString
    Negative distance values are taken as measured in the reverse
    direction from the end of the geometry. Out-of-range index
    values are handled by clamping them to the valid range of values.
    If the start distance equals the end distance, a Point is returned.
    If the start distance is actually beyond the end distance, then the
    reversed substring is returned such that the start distance is
    at the first coordinate.
    Parameters
    ----------
    geom : LineString
        The geometry to get a substring of.
    start_dist : float
        The distance along `geom` of the start of the substring.
    end_dist : float
        The distance along `geom` of the end of the substring.
    normalized : bool, False
        Whether the distance parameters are interpreted as a
        fraction of the geometry's length.
    Returns
    -------
    Union[Point, LineString]
        The substring between `start_dist` and `end_dist` or a Point
        if they are at the same location.
    Raises
    ------
    TypeError
        If `geom` is not a LineString.
    Examples
    --------
    >>> ls = LineString((i, 0) for i in range(6))
    >>> ls.wkt
    'LINESTRING (0 0, 1 0, 2 0, 3 0, 4 0, 5 0)'
    >>> substring(ls, start_dist=1, end_dist=3).wkt
    'LINESTRING (1 0, 2 0, 3 0)'
    >>> substring(ls, start_dist=3, end_dist=1).wkt
    'LINESTRING (3 0, 2 0, 1 0)'
    >>> substring(ls, start_dist=1, end_dist=-3).wkt
    'LINESTRING (1 0, 2 0)'
    >>> substring(ls, start_dist=0.2, end_dist=-0.6, normalized=True).wkt
    'LINESTRING (1 0, 2 0)'
    Returning a `Point` when `start_dist` and `end_dist` are at the
    same location.
    >>> substring(ls, 2.5, -2.5).wkt
    'POINT (2.5 0)'
    """

    if not isinstance(geom, LineString):
        raise TypeError("Can only calculate a substring of LineString geometries. A %s was provided." % geom.type)

    # Filter out cases in which to return a point
    if start_dist == end_dist:
        return geom.interpolate(start_dist, normalized)
    elif not normalized and start_dist >= geom.length and end_dist >= geom.length:
        return geom.interpolate(geom.length, normalized)
    elif not normalized and -start_dist >= geom.length and -end_dist >= geom.length:
        return geom.interpolate(0, normalized)
    elif normalized and start_dist >= 1 and end_dist >= 1:
        return geom.interpolate(1, normalized)
    elif normalized and -start_dist >= 1 and -end_dist >= 1:
        return geom.interpolate(0, normalized)

    if normalized:
        start_dist *= geom.length
        end_dist *= geom.length

    # Filter out cases where distances meet at a middle point from opposite ends.
    if start_dist < 0 < end_dist and abs(start_dist) + end_dist == geom.length:
        return geom.interpolate(end_dist)
    elif end_dist < 0 < start_dist and abs(end_dist) + start_dist == geom.length:
        return geom.interpolate(start_dist)

    start_point = geom.interpolate(start_dist)
    end_point = geom.interpolate(end_dist)

    if start_dist < 0:
        start_dist = geom.length + start_dist  # Values may still be negative,
    if end_dist < 0:                           # but only in the out-of-range
        end_dist = geom.length + end_dist      # sense, not the wrap-around sense.

    reverse = start_dist > end_dist
    if reverse:
        start_dist, end_dist = end_dist, start_dist

    if start_dist < 0:
        start_dist = 0  # to avoid duplicating the first vertex

    if reverse:
        vertex_list = [(end_point.x, end_point.y)]
    else:
        vertex_list = [(start_point.x, start_point.y)]

    coords = list(geom.coords)
    current_distance = 0
    for p1, p2 in zip(coords, coords[1:]):
        if start_dist < current_distance < end_dist:
            vertex_list.append(p1)
        elif current_distance >= end_dist:
            break

        current_distance += ((p2[0] - p1[0]) ** 2 + (p2[1] - p1[1]) ** 2) ** 0.5

    if reverse:
        vertex_list.append((start_point.x, start_point.y))
        # reverse direction result
        vertex_list = reversed(vertex_list)
    else:
        vertex_list.append((end_point.x, end_point.y))

    return LineString(vertex_list)


def reverse_linestring(linestring):
    coords_list = list(linestring.coords)
    coords_list.reverse()
    return LineString(coords_list)


def rubber_snap(pt, linestring, fix_st):
    """
    捕捉

    Args:
        pt: 目标点
        linestring: 进行捕捉的线段
        fix_st: 是否固定linestring的起点。若fix_st=True，则用linestring的终点捕捉到pt；反之则用起点捕捉到pt。

    Returns:

    """
    if not fix_st:
        linestring = reverse_linestring(linestring)

    dxn = pt.coords[0][0] - linestring.coords[-1][0]
    dyn = pt.coords[0][1] - linestring.coords[-1][1]
    point_num = len(linestring.coords)

    new_xy = []
    for index, pxy in enumerate(list(linestring.coords)):
        px = dxn * index / (point_num - 1) + pxy[0]
        py = dyn * index / (point_num - 1) + pxy[1]
        new_xy.append((px, py))
    new_geom = LineString(new_xy)
    if not fix_st:
        new_geom = reverse_linestring(new_geom)
    return new_geom


def rubber_snap_within_distance(pt_target, linestring, fix_st, distance):
    """
    将指定端点附近distance长度以内的部分进行橡皮变换捕捉到指定点(pt_target)，其余部分不变。

    Args:
        pt_target:
        linestring:
        fix_st:
        distance:

    Returns:

    """
    if not fix_st:
        linestring = reverse_linestring(linestring)
    else:
        linestring = linestring
    affine_len = 0
    coords_list = list(linestring.coords)
    fix_index = 0
    for fix_index in range(len(coords_list) - 2, 0, -1):
        affine_len += cal_dist(coords_list[fix_index], coords_list[fix_index + 1])
        if affine_len > distance:
            break

    # 橡皮变换
    affine_geom = rubber_snap(pt_target, LineString(coords_list[fix_index:]), fix_st=True)
    new_geom = LineString(coords_list[:fix_index] + list(affine_geom.coords))
    if not fix_st:
        new_geom = reverse_linestring(new_geom)
    return new_geom


def compare_road_rank(r_class1, r_class2):
    """
    比较r_class1和r_class2的级别。若r_class1级别高于r_class2，返回1，相等返回0，低于返回-1

    Args:
        r_class1:
        r_class2:

    Returns: boolean

    """
    if ROAD_RANK[r_class1] < ROAD_RANK[r_class2]:  # 级别高，数字小
        return 1
    elif ROAD_RANK[r_class1] == ROAD_RANK[r_class2]:  # 级别相同
        return 0
    else:
        return -1
